package cn.iocoder.boot.kmeans;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class KMeans {
    public static List<Cluster> cluster(List<Talent> talents, int k, int maxIterations) {
        List<Cluster> clusters = new ArrayList<>();
        Random random = new Random();
        for (int i = 0; i < k; i++) {
            Talent initialCentroid = new Talent("Centroid", talents.get(random.nextInt(talents.size())).getFeatures(), 0);
            Cluster cluster = new Cluster();
            cluster.setCentroid(initialCentroid);
            clusters.add(cluster);
        }

        for (int iteration = 0; iteration < maxIterations; iteration++) {
            for (Cluster cluster : clusters) {
                cluster.getTalents().clear();
            }

            for (Talent talent : talents) {
                Cluster nearestCluster = null;
                double minDistance = Double.MAX_VALUE;
                for (Cluster cluster : clusters) {
                    double distance = talent.distanceTo(cluster.getCentroid());
                    if (distance < minDistance) {
                        minDistance = distance;
                        nearestCluster = cluster;
                    }
                }
                nearestCluster.addTalent(talent);
            }

            boolean converged = true;
            for (Cluster cluster : clusters) {
                Cluster oldCluster = new Cluster();
                oldCluster.setCentroid(cluster.getCentroid());
                oldCluster.getTalents().addAll(cluster.getTalents());
                cluster.updateCentroid();
                if (!isSameCentroid(oldCluster.getCentroid(), cluster.getCentroid())) {
                    converged = false;
                }
            }

            if (converged) {
                break;
            }
        }

        return clusters;
    }

    private static boolean isSameCentroid(Talent centroid1, Talent centroid2) {
        if (centroid1 == null || centroid2 == null) {
            return false;
        }
        if (centroid1.getFeatures().length != centroid2.getFeatures().length) {
            return false;
        }
        for (int i = 0; i < centroid1.getFeatures().length; i++) {
            if (centroid1.getFeatures()[i] != centroid2.getFeatures()[i]) {
                return false;
            }
        }
        return true;
    }
}